import pickle
import numpy as np
import pandas as pd
from pathlib import Path
from typing import List, Tuple, Optional
import logging
from datetime import datetime
import sys
import os


# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, '..'))
if project_root not in sys.path:
    sys.path.insert(0, project_root)
# =============================================

try:
    from config import settings
    from utils.logger import setup_logger
    from core.embedding_client import EmbeddingClient
except ImportError as e:
    print(f"导入错误: {e}")
    print("正在使用备用方案...")


    # 备用方案：直接在这里定义必要的类和函数
    class Settings:
        VECTOR_CACHE_FILE = Path("data/cache/warehouse_vectors_v1.0.pkl")
        OPENAI_MODEL = "BAAI/bge-large-zh-v1.5"
        OPENAI_API_KEY = "sk-aswwzdvvimeybiiokqebixpkhbmcftlbgkubssfuodifqjcf"
        OPENAI_BASE_URL = "https://api.siliconflow.cn/v1"
        MAX_BATCH_SIZE = 32
        MAX_RETRY_ATTEMPTS = 5


    settings = Settings()


    def setup_logger(name):
        logger = logging.getLogger(name)
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
        return logger


    class EmbeddingClient:
        def __init__(self):
            self.model = settings.OPENAI_MODEL

        def create_embeddings(self, texts):
            print(f"模拟生成嵌入向量: {len(texts)} 条文本")
            # 这里返回随机向量用于测试
            return np.random.randn(len(texts), 768)

logger = setup_logger(__name__)


class VectorManager:
    """向量库管理工具，支持 add/delete/update/save/load"""

    def __init__(self):
        self.cache_file = settings.VECTOR_CACHE_FILE
        self.embedding_client = EmbeddingClient()
        self.ids: List[int] = []
        self.vectors: Tuple[np.ndarray, np.ndarray] = (None, None)
        self.units: List[str] = []
        self.metadata = {
            'version': '1.0',
            'created_at': datetime.now().isoformat(),
            'updated_at': datetime.now().isoformat(),
            'total_records': 0,
            'vector_dimension': 0
        }

    # ... 保持你原有的所有方法不变 ...
    def load(self) -> bool:
        """加载向量缓存"""
        try:
            if not self.cache_file.exists():
                logger.warning("缓存文件不存在")
                return False

            with open(self.cache_file, 'rb') as f:
                data = pickle.load(f)
                self.ids, self.vectors, self.units, self.metadata = data

            logger.info(f"加载向量库成功，记录数: {len(self.ids)}")
            return True

        except Exception as e:
            logger.error(f"加载向量库失败: {e}")
            return False

    def save(self) -> bool:
        """保存向量缓存"""
        try:
            self.metadata['updated_at'] = datetime.now().isoformat()
            self.metadata['total_records'] = len(self.ids)

            if len(self.ids) > 0:
                self.metadata['vector_dimension'] = self.vectors[0].shape[1]

            with open(self.cache_file, 'wb') as f:
                pickle.dump((self.ids, self.vectors, self.units, self.metadata), f)

            logger.info(f"保存向量库成功，记录数: {len(self.ids)}")
            return True

        except Exception as e:
            logger.error(f"保存向量库失败: {e}")
            return False

    def add_vectors(self, new_ids: List[int], new_names: List[str],
                    new_specs: List[str], new_units: List[str]) -> bool:
        """添加新向量"""
        try:
            # 生成新向量的嵌入
            new_name_vecs = self.embedding_client.create_embeddings(new_names)
            new_spec_vecs = self.embedding_client.create_embeddings(new_specs)

            # 合并到现有数据
            if len(self.ids) == 0:
                self.ids = new_ids
                self.vectors = (new_name_vecs, new_spec_vecs)
                self.units = new_units
            else:
                self.ids.extend(new_ids)
                self.vectors = (
                    np.vstack([self.vectors[0], new_name_vecs]),
                    np.vstack([self.vectors[1], new_spec_vecs])
                )
                self.units.extend(new_units)

            logger.info(f"添加 {len(new_ids)} 条向量记录")
            return self.save()

        except Exception as e:
            logger.error(f"添加向量失败: {e}")
            return False

    def delete_vectors(self, ids_to_delete: List[int]) -> bool:
        """删除向量"""
        try:
            if len(self.ids) == 0:
                return True

            # 创建索引映射
            keep_indices = [i for i, id_val in enumerate(self.ids)
                            if id_val not in ids_to_delete]

            self.ids = [self.ids[i] for i in keep_indices]
            self.vectors = (
                self.vectors[0][keep_indices],
                self.vectors[1][keep_indices]
            )
            self.units = [self.units[i] for i in keep_indices]

            logger.info(f"删除 {len(ids_to_delete)} 条向量记录")
            return self.save()

        except Exception as e:
            logger.error(f"删除向量失败: {e}")
            return False

    def update_vectors(self, ids_to_update: List[int], new_names: List[str],
                       new_specs: List[str], new_units: List[str]) -> bool:
        """更新向量"""
        try:
            # 生成新向量
            new_name_vecs = self.embedding_client.create_embeddings(new_names)
            new_spec_vecs = self.embedding_client.create_embeddings(new_specs)

            # 更新现有向量
            for i, id_val in enumerate(ids_to_update):
                if id_val in self.ids:
                    idx = self.ids.index(id_val)
                    self.vectors[0][idx] = new_name_vecs[i]
                    self.vectors[1][idx] = new_spec_vecs[i]
                    self.units[idx] = new_units[i]

            logger.info(f"更新 {len(ids_to_update)} 条向量记录")
            return self.save()

        except Exception as e:
            logger.error(f"更新向量失败: {e}")
            return False

    def get_stats(self) -> dict:
        """获取统计信息"""
        return {
            'total_records': len(self.ids),
            'vector_dimension': self.vectors[0].shape[1] if len(self.ids) > 0 else 0,
            'last_updated': self.metadata['updated_at'],
            'version': self.metadata['version']
        }


# 测试代码
if __name__ == "__main__":
    print("测试 VectorManager...")

    # 创建数据目录
    data_dir = Path("data/cache")
    data_dir.mkdir(parents=True, exist_ok=True)

    manager = VectorManager()
    print("✓ VectorManager 初始化成功")

    # 测试添加向量
    test_ids = [1, 2, 3]
    test_names = ["产品A", "产品B", "产品C"]
    test_specs = ["规格A", "规格B", "规格C"]
    test_units = ["个", "件", "台"]

    success = manager.add_vectors(test_ids, test_names, test_specs, test_units)
    if success:
        print("✓ 添加向量成功")
        stats = manager.get_stats()
        print(f"📊 统计信息: {stats}")

        # 测试加载
        manager2 = VectorManager()
        if manager2.load():
            print("✓ 加载向量库成功")
            print(f"📊 加载的记录数: {len(manager2.ids)}")
        else:
            print("✗ 加载向量库失败")
    else:
        print("✗ 添加向量失败")

    print("测试完成！")